import csv
import random
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from util.util_image import draw_img_xyz, draw_img
from util.util_csv import read_csv_to_numpy_array
from util.util_ris_pattern_2 import phase_2_pattern_xyz, eps


# ============================================ 工具代码 ================================================
# 统一画图
def draw_img_xyz_paper(data, x, y):
    # 绘制图像
    plt.figure(figsize=(10, 8))
    plt.grid(True)
    # plt.pcolormesh(x, y, data, shading='auto', cmap='plasma')
    plt.pcolormesh(x, y, data, shading='auto', cmap='gist_rainbow')
    plt.colorbar(label='Pattern Magnitude')
    plt.axis('equal')
    # plt.axis('off')  # 关闭坐标轴显示
    plt.show()


def draw_img_xyz_paper_circular(data, x, y, center=(0, 0), radius=1.0):
    """
    绘制一个二维热力图，但只显示指定圆形区域内的数据。

    参数:
    data : 2D numpy array
        要绘制的数据。
    x : 1D or 2D numpy array
        X 坐标。
    y : 1D or 2D numpy array
        Y 坐标。
    center : tuple, optional
        圆的中心坐标 (x0, y0)。默认是 (0, 0)。
    radius : float, optional
        圆的半径。默认是 1.0。
    """
    # 确保x, y是网格形式
    if x.ndim == 1 and y.ndim == 1:
        X, Y = np.meshgrid(x, y)
    elif x.ndim == 2 and y.ndim == 2:
        X, Y = x, y
    else:
        raise ValueError("x和y必须是1D或2D数组，并且形状要与data兼容。")

    # 1. 创建圆形遮罩
    # 计算每个网格点到指定中心的距离的平方
    distance_squared = (X - center[0]) ** 2 + (Y - center[1]) ** 2
    # 创建遮罩：距离平方大于半径平方的点为True (需要被遮蔽)
    mask = distance_squared > radius ** 2

    # 2. 将圆外的数据设置为 NaN
    # np.ma.masked_where 会创建一个 MaskedArray，被遮罩为True的位置会被隐藏
    masked_data = np.ma.masked_where(mask, data)

    # 3. 绘制图像
    plt.figure(figsize=(10, 8))
    plt.grid(True)

    # 使用 MaskedArray 进行绘制，圆外区域不会被显示
    mesh = plt.pcolormesh(X, Y, masked_data, shading='auto', cmap='gist_rainbow')

    plt.colorbar(mesh, label='Pattern Magnitude')
    plt.axis('equal')
    # 可选：设置坐标轴范围以聚焦在圆形区域附近
    # plt.xlim(center[0] - radius * 1.1, center[0] + radius * 1.1)
    # plt.ylim(center[1] - radius * 1.1, center[1] + radius * 1.1)
    plt.title(f'Heatmap (Circular Region: Center={center}, Radius={radius})')
    plt.show()


# 读取3d的.csv
def read_from_csv(path_csv):
    x_coords = []
    y_coords = []
    z_values = []

    with open(path_csv, 'r') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            x_coords.append(float(row['X']))
            y_coords.append(float(row['Y']))
            z_values.append(float(row['Z']))

    x_coords = np.unique(np.array(x_coords))
    y_coords = np.unique(np.array(y_coords))
    z_values = np.array(z_values).reshape((len(x_coords), len(y_coords)))

    return x_coords, y_coords, z_values


def read_from_csv_2(path_csv):
    """
    读取CSV文件，并将'Z'值填入一个181x181的预初始化数组中。

    初始化的数组尺寸为181x181，每个单元的初始值是-74.3314和-61.405之间的随机小数。
    CSV中的('X', 'Y', 'Z')数据会被映射到数组的坐标('X'+90, 'Y'+90)处。

    参数:
        path_csv (str): CSV文件的路径。

    返回:
        tuple: (x_coords, y_coords, z_matrix)
            x_coords (numpy.ndarray): 从-90到90的X坐标数组 (用于参考)。
            y_coords (numpy.ndarray): 从-90到90的Y坐标数组 (用于参考)。
            z_matrix (numpy.ndarray): 填充后的181x181二维数组。
    """
    # 1. 初始化 181x181 的二维数组，填充随机小数
    # np.random.uniform 更高效地生成随机数数组
    z_matrix = np.random.uniform(low=-74.3314, high=-61.405, size=(181, 181))

    # 2. 读取 CSV 文件并填充数组
    with open(path_csv, 'r') as csvfile:
        # 假设CSV有标题行 'X', 'Y', 'Z'
        reader = csv.DictReader(csvfile)
        for row in reader:
            try:
                x = float(row['X'])
                y = float(row['Y'])
                z = float(row['Z'])

                # 3. 计算映射到 z_matrix 的索引
                # 根据规则：数组索引 = 原始坐标 + 90
                # 由于数组索引是整数，我们需要将浮点坐标四舍五入或转换为整数。
                # 假设原始 X, Y 坐标是整数（例如 -90 到 90）
                i = int(round(x + 90))
                j = int(round(y + 90))

                # 4. 检查索引是否在有效范围内
                if 0 <= i < 181 and 0 <= j < 181:
                    z_matrix[i, j] = z
                else:
                    print(f"警告：坐标 ({x}, {y}) 映射到的索引 ({i}, {j}) 超出数组范围，已忽略。")

            except (ValueError, KeyError) as e:
                print(f"警告：处理行时出错 {row}: {e}")

    # 5. 创建用于参考的 x_coords 和 y_coords 数组
    # 这些数组代表了 z_matrix 每个索引对应的原始坐标值
    x_coords = np.arange(-90, 91)  # -90 to 90 inclusive
    y_coords = np.arange(-90, 91)  # -90 to 90 inclusive

    return x_coords, y_coords, z_matrix

# ============================================ 业务代码 -- 方向图 python ================================================
# 画2d俯视方向图 -- Python模拟
def plot_pattern2d_python(path_csv_phaseBit):
    # 读取码阵
    phaseBit = read_csv_to_numpy_array(path_csv_phaseBit)
    # 计算方向图
    pattern_xyz, x, y, z = phase_2_pattern_xyz(phaseBit)
    # pattern_xyz = np.abs(pattern_xyz)
    pattern_dbw_xyz = 20 * np.log10(np.abs(pattern_xyz) / np.max(np.max(np.abs(pattern_xyz))) + eps)
    indices = (pattern_dbw_xyz < -40).nonzero()
    pattern_dbw_xyz[indices] = -40
    # 画2d验证
    # draw_img_xyz(pattern_dbw_xyz, x, y)
    # 画2d俯视方向图
    draw_img_xyz_paper(pattern_dbw_xyz, x, y)


# ============================================ 业务代码 -- 方向图 HFSS ================================================
# HFSS方向图转成360*360矩阵
def format_hfss3d_2_360x360(path_hfss3d):
    pattern = np.zeros((360, 90))
    pattern_line = np.array(pd.read_csv(path_hfss3d))

    for ii in pattern_line:
        if ii[0] >= 0:
            if ii[0] < 360 and ii[1] < 90:
                pattern[int(ii[0])][int(ii[1])] = ii[2]
        else:
            pattern[int(360 + ii[0])][int(ii[1])] = ii[2]

    return pattern


# 画2d俯视方向图 -- HFSS模拟
def plot_pattern2d_hfss(path_hfss3d):
    pattern = format_hfss3d_2_360x360(path_hfss3d)
    #
    theta = np.linspace(0, np.pi / 2, 90)
    phi = np.linspace(0, 2 * np.pi * (1 - 1 / 360), 360)
    th, ph = np.meshgrid(theta, phi)
    # 2. 转换为直角坐标系
    x = np.sin(th) * np.cos(ph)
    y = np.sin(th) * np.sin(ph)
    z = np.cos(th)
    # 转dB
    pattern_dbw = 20 * np.log10(np.abs(pattern) / np.max(np.max(np.abs(pattern))) + eps)
    indices = (pattern_dbw < -40).nonzero()
    pattern_dbw[indices] = -40
    # 画个只有pattern的确认下
    # draw_img(pattern_dbw)
    # 画个带xy的pattern确认下
    # draw_img_xyz(pattern_dbw, x, y)
    # 画2d俯视方向图
    draw_img_xyz_paper(pattern_dbw, x, y)


# ============================================ 业务代码 -- 方向图 实测 ================================================
# 画2d俯视方向图 -- 实测模拟
def plot_pattern2d_measurement(path_measure3d):
    th_coords, ph_coords, pattern_matrix = read_from_csv_2(path_measure3d)
    pattern_matrix = np.exp(pattern_matrix / 20)
    pattern_matrix = pattern_matrix / np.max(pattern_matrix)
    pattern_matrix = 20 * np.log10(pattern_matrix)
    indices = (pattern_matrix < -40).nonzero()
    pattern_matrix[indices] = -40
    #
    th_coords = np.deg2rad(th_coords)
    ph_coords = np.deg2rad(ph_coords)
    x_coords = np.sin(th_coords)
    y_coords = np.sin(ph_coords)
    # 画个只有pattern的确认下
    # draw_img(pattern_matrix)
    # 画个带xy的pattern确认下
    # draw_img_xyz(pattern_matrix, x_coords, y_coords)
    # 画2d俯视方向图
    # draw_img_xyz_paper(pattern_matrix, x_coords, y_coords)
    draw_img_xyz_paper_circular(pattern_matrix, x_coords, y_coords, center=(0, 0), radius=1.0)


# 画2d俯视方向图 -- 实测模拟 (数据直接减)
def plot_pattern2d_measurement_2(path_measure3d):
    th_coords, ph_coords, pattern_matrix = read_from_csv_2(path_measure3d)
    mmm = np.max(pattern_matrix)
    print(f"mmm: {mmm}")
    pattern_matrix = pattern_matrix - np.max(pattern_matrix)
    indices = (pattern_matrix < -40).nonzero()
    pattern_matrix[indices] = -40
    #
    th_coords = np.deg2rad(th_coords)
    ph_coords = np.deg2rad(ph_coords)
    x_coords = np.sin(th_coords)
    y_coords = np.sin(ph_coords)
    # 画个只有pattern的确认下
    # draw_img(pattern_matrix)
    # 画个带xy的pattern确认下
    # draw_img_xyz(pattern_matrix, x_coords, y_coords)
    # 画2d俯视方向图
    # draw_img_xyz_paper(pattern_matrix, x_coords, y_coords)
    draw_img_xyz_paper_circular(pattern_matrix, x_coords, y_coords, center=(0, 0), radius=1.0)



if __name__ == "__main__":
    # 画2d俯视方向图 -- python
    # plot_pattern2d_python(
    #     "../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/theta(0,60)-phi(0,0)-20250816/phaseBit_AGA_(0,0).csv")
    # plot_pattern2d_python(
    #     "../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/beam2-(20,45)-(20,225)/phaseBit_AGA_(((20,45),(20,225))).csv")
    # 画2d俯视方向图 -- HFSS
    # plot_pattern2d_hfss(
    #     "../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/32x32-beam1(0,0)-rE_Plot_3d.csv")
    # plot_pattern2d_hfss(
    #     "../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/32x32-beam2(20,45)(20,225)-rE_Plot_3d.csv")
    # 画2d俯视方向图 -- 实测
    # plot_pattern2d_measurement("../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/1-3D.csv")
    # plot_pattern2d_measurement_2("../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/1-3D.csv")
    # plot_pattern2d_measurement("../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/10-3D.csv")
    # plot_pattern2d_measurement_2("../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/10-3D.csv")
    #
    #
    # 画单波束扫描方向图
    # plot_pattern2d_python(
    #     "../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/theta(45,45)-phi(0,0)-AGAiteration100/phaseBit_AGA_(45,0).csv")
    plot_pattern2d_python(
        "../../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/beam2-(20,45)-(20,135)/phaseBit_AGA_(((20,45),(20,135))).csv")
